Add byte-progress reporting for safetensors file loading - #427
Open
aleroot wants to merge 2 commits into
Open
Conversation
Contributor
Author
|
This now depends on ml-explore/mlx#3742 , I will update this PR when that one is going to be merged. |
davidkoski
reviewed
Jun 30, 2026
Comment on lines
+9
to
+10
| /// single load are delivered in monotonically increasing order. | ||
| public struct LoadProgress: Sendable, Equatable { |
Member
There was a problem hiding this comment.
Is it worth documenting that:
- it may only read part of the file -- if safetensors are not consumed then it won't read them
- I don't know if it is possible, but it could read part of the file twice and it could go over 100%
Not an issue but might be surprising (loading safetensors is lazy but not in an obvious way).
davidkoski
reviewed
Jun 30, 2026
Member
There was a problem hiding this comment.
I can't find the 00375d0 commit in mlx-c -- I presume this PR requires a newer mlx/mlx-c?
Contributor
Author
There was a problem hiding this comment.
Yes, before we can merge this, this other one need to be merged first, then this one needs to be updated.
Contributor
Author
There was a problem hiding this comment.
Today a new version of mlx has been released 🍾 , opened this one ml-explore/mlx-c#126
aleroot
force-pushed
the
loading_progress
branch
from
August 18, 2026 14:56
9939e79 to
5140312
Compare
Introduces LoadProgress and new loadArrays(url:stream:progressHandler:) and loadArraysAndMetadata(url:stream:progressHandler:) overloads that report bytes read as lazy arrays are evaluated. Uses a custom mlx_io_reader vtable backed by pread() so progress callbacks can be invoked from MLX worker threads. Includes a unit test verifying monotonic progress from 0 to 1.
`loadWeights()` style helpers -- including the one in mlx-swift-lm -- call the
plain `loadArrays(url:)` / `loadArraysAndMetadata(url:)`, so a per-call
`progressHandler:` argument can never reach them without changing every caller
along the way.
Add `withLoadProgressHandler(_:_:)` (sync and async), a task local scoped
handler in the style of `withErrorHandler(_:_:)`. The plain file loading
functions report byte progress to it when one is installed, so an application
can drive a precise model loading progress bar around code it does not own:
let container = try await withLoadProgressHandler({ tracker.update($0) }) {
try await factory.loadContainer(from: directory, using: tokenizerLoader)
}
`LoadProgress` gains the `url` of the file being read so progress can be
aggregated across the shards of a sharded model.
Also fix the SEEK_END case of the in-memory reader, which moved the offset
relative to the current position instead of the end of the data. mlx 0.32.1
seeks to the end of the stream to validate the tensor data offsets against the
size of the file, so `loadArrays(data:)` would fail there ("The JSON header is
N bytes long but the file is only 8 bytes").
Finally, restructure the truncated file test: a truncated file may now be
reported either eagerly, while the header is parsed, or lazily, when the arrays
are evaluated, and neither happens with the currently vendored mlx/mlx-c.
aleroot
force-pushed
the
loading_progress
branch
from
August 18, 2026 15:03
5140312 to
fdb7f7a
Compare
aleroot
added a commit
to aleroot/mlx-swift-lm
that referenced
this pull request
Aug 18, 2026
Loading a model is dominated by reading the weights from disk, but nothing is
reported while that happens: `progressHandler` only covers the download, so an
application showing "loading..." has no way to draw an accurate progress bar,
and a large model can spend tens of seconds there.
Group the progress callbacks of a load in one value, `LoadProgressHandlers`:
let container = try await LLMModelFactory.shared.loadContainer(
from: directory, using: tokenizerLoader,
progress: .weights { progress in
print(progress.fractionCompleted)
})
It is built on the scoped progress handler of mlx-swift, so `_load()` and the
model implementations are untouched -- the handler is installed around the load
and the plain `loadArrays(url:)` calls in `loadWeights()` report to it.
`ModelLoadProgressReporter` aggregates the byte progress that MLX reports per
file -- a model is frequently split into several shards, read concurrently --
into a single `Progress` for the whole model, and coalesces the updates, as MLX
reports roughly one per 4MB and the handler typically hops to the main actor.
Loading is lazy, so the weights are read while the model is evaluated at the end
of `loadWeights()`. The weights that `sanitize(weights:metadata:)` drops are
never evaluated, and therefore never read, so the aggregate can legitimately
stop short of the size of the files: completion is published once the load
returns.
The existing `progressHandler` parameter is unchanged and keeps reporting the
download; `progress.download` is called in addition to it when both are given.
Note: requires the scoped `withLoadProgressHandler(_:_:)` of
ml-explore/mlx-swift#427.
aleroot
added a commit
to aleroot/mlx-swift-lm
that referenced
this pull request
Aug 18, 2026
Loading a model is dominated by reading the weights from disk, but nothing is
reported while that happens: `progressHandler` only covers the download, so an
application showing "loading..." has no way to draw an accurate progress bar,
and a large model can spend tens of seconds there.
Group the progress callbacks of a load in one value, `LoadProgressHandlers`:
let container = try await LLMModelFactory.shared.loadContainer(
from: directory, using: tokenizerLoader,
progress: .weights { progress in
print(progress.fractionCompleted)
})
It is built on the scoped progress handler of mlx-swift, so `_load()` and the
model implementations are untouched -- the handler is installed around the load
and the plain `loadArrays(url:)` calls in `loadWeights()` report to it.
`ModelLoadProgressReporter` aggregates the byte progress that MLX reports per
file -- a model is frequently split into several shards, read concurrently --
into a single `Progress` for the whole model, and coalesces the updates, as MLX
reports roughly one per 4MB and the handler typically hops to the main actor.
Loading is lazy, so the weights are read while the model is evaluated at the end
of `loadWeights()`. The weights that `sanitize(weights:metadata:)` drops are
never evaluated, and therefore never read, so the aggregate can legitimately
stop short of the size of the files: completion is published once the load
returns.
The existing `progressHandler` parameter is unchanged and keeps reporting the
download; `progress.download` is called in addition to it when both are given.
Note: requires the scoped `withLoadProgressHandler(_:_:)` of
ml-explore/mlx-swift#427.
aleroot
added a commit
to aleroot/mlx-swift-lm
that referenced
this pull request
Aug 18, 2026
Loading a model is dominated by reading the weights from disk, but nothing is
reported while that happens: `progressHandler` only covers the download, so an
application showing "loading..." has no way to draw an accurate progress bar,
and a large model can spend tens of seconds there.
Group the progress callbacks of a load in one value, `LoadProgressHandlers`:
let container = try await LLMModelFactory.shared.loadContainer(
from: directory, using: tokenizerLoader,
progress: .weights { progress in
print(progress.fractionCompleted)
})
It is built on the scoped progress handler of mlx-swift, so `_load()` and the
model implementations are untouched -- the handler is installed around the load
and the plain `loadArrays(url:)` calls in `loadWeights()` report to it.
`ModelLoadProgressReporter` aggregates the byte progress that MLX reports per
file -- a model is frequently split into several shards, read concurrently --
into a single `Progress` for the whole model, and coalesces the updates, as MLX
reports roughly one per 4MB and the handler typically hops to the main actor.
Loading is lazy, so the weights are read while the model is evaluated at the end
of `loadWeights()`. The weights that `sanitize(weights:metadata:)` drops are
never evaluated, and therefore never read, so the aggregate can legitimately
stop short of the size of the files: completion is published once the load
returns.
The existing `progressHandler` parameter is unchanged and keeps reporting the
download; `progress.download` is called in addition to it when both are given.
Note: requires the scoped `withLoadProgressHandler(_:_:)` of
ml-explore/mlx-swift#427.
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Proposed changes
Adds byte-level progress reporting for safetensors file loading so upper layers can drive precise model-loading progress UI.
This PR depends on ml-explore/mlx#3734.
The progress API in mlx-swift can report byte-level safetensors loading progress, but the truncated/failed-read correctness path needs the upstream MLX fix first. Without ml-explore/mlx#3734, CPU lazy-load read failures can be swallowed before they reach Swift.
Once ml-explore/mlx#3734 eventaully lands, I’ll update the MLX submodule reference here and rerun the mlx-swift test suite.
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes